module FES3DT412CTE_m
use Mesh3DT4_m
use FES3DT4_m
use FES3DT41_m
use FES3DT42_m
use Geometry_m
implicit none

type, extends(FES3DT4) :: FES3DT412CTE
contains
    procedure :: initFES => initFES3DT412CTE
    
    procedure :: assembleHeatFluxT4CTE_Const
    procedure :: assembleHeatFluxT4CTE_function
    generic :: assembleHeatFlux => assembleHeatFluxT4CTE_Const, assembleHeatFluxT4CTE_function
    
    procedure :: assembleHeatConvectSurfaceT4CTE_const
    generic :: assembleHeatConvectSurface => assembleHeatConvectSurfaceT4CTE_const
    
    procedure :: assembleDC1 => assembleDC1_T4CTE
    procedure :: assembleDC2 => assembleDC2_T4CTE
    
    procedure :: applyBCHeatT4CTE_function
    procedure :: applyBCHeatT4CTE_const
    generic :: applyBCHeat => applyBCHeatT4CTE_function, applyBCHeatT4CTE_const
    
    procedure :: applyBCZeroDisp => applyBCZeroDispT4CTE
    procedure :: assembleHeatToElastic => assembleHeatToElasticT4CTE
    procedure :: extractUVWH => extractUVWHT4CTE
end type

contains
subroutine initFES3DT412CTE(this, mesh3D)
    implicit none
    class(FES3DT412CTE), intent(out) :: this
    type(Mesh3DT4), target, intent(in) :: mesh3D

    integer :: band
    this%dim = 3
    this%nDof = mesh3D%nv * 4
    this%nVertices = mesh3D%nv
    this%nElement = mesh3D%nt
    this%nBoundary = mesh3D%nb
    this%nPerElement = 4
    this%nPerBoundary = 3
    this%dofPerNode = 4
    this%dofPerElement = 16
    this%dofPerBoundary = 12
    
    band = mesh3D.elementDiff()
    this%nBand = (band + 1) * 4
    this%mesh3D => mesh3D
end subroutine

subroutine assembleHeatToElasticT4CTE(this, beta, tpCoupled, A)
    implicit none
    class(FES3DT412CTE), intent(in) :: this
    real*8, intent(in) :: beta
    integer, intent(in) :: tpCoupled
    type(SMatrix) :: A

    integer :: i, j, idx(4), idx2(12)
    real*8 :: xyz(4, 3), bbeta(4), HeatToElastic(12, 4)
    type(Mesh3DT4), pointer :: mesh3D
    mesh3D => this%mesh3D
    bbeta = beta
    do i = 1, this%nElement
        idx = mesh3D%elements(i)%idx
        xyz(:, 1) = mesh3D%nodes(idx)%xyz(1)
        xyz(:, 2) = mesh3D%nodes(idx)%xyz(2)
        xyz(:, 3) = mesh3D%nodes(idx)%xyz(3)
        call getDofFromElement(idx, 3, idx2)
        call elementCoupledMatrixT4CTE(xyz, bbeta, tpCoupled, HeatToElastic)
        call assembleMatrix(A, HeatToElastic, idx2, idx)
    end do
end subroutine

subroutine assembleDC1_T4CTE(this, r, c, k, TRef, E, Nu, beta, dt, &
                             tpElasticStiff, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES3DT412CTE), intent(in) :: this
    real*8, intent(in):: r, c, k, TRef, E, Nu, beta, dt
    integer, intent(in) :: tpElasticStiff, tpHeatStiff, tpHeatMass, tpCoupled
    type(SMatrix), intent(inout) :: LHS, RHS    

    integer :: i, j, idx(4), idxHeat(4), idxElastic(12)
    real*8 :: xyz(4, 3), HeatStiff(4, 4), HeatMass(4, 4)
    real*8 :: coefM(4), coefK(4), coefE(4), coefNu(4), coefBeta(4)
    real*8 :: ElasticStiff(12, 12), HeatToElastic(12, 4), ElasticToHeat(4, 12)
    type(Mesh3DT4), pointer :: mesh3D
    mesh3D => this%mesh3D
    coefM = r * c / TRef
    coefK = k / TRef
    coefE = E
    coefNu = Nu
    coefBeta = beta
    do i = 1, mesh3D%nt
        idx = mesh3D%elements(i)%idx
        xyz(:, 1) = mesh3D%nodes(idx)%xyz(1)
        xyz(:, 2) = mesh3D%nodes(idx)%xyz(2)
        xyz(:, 3) = mesh3D%nodes(idx)%xyz(3)
        call getDofFromElement(idx, 3, idxElastic)
        idxHeat = idx + 3 * mesh3D%nv
        call elementElasticityStiff3DT42(xyz, coefE, coefNu, tpElasticStiff, ElasticStiff)
        call assembleMatrix(LHS, ElasticStiff, idxElastic)
        
        call elementCoupledMatrixT4CTE(xyz, coefBeta, tpCoupled, HeatToElastic)
        ElasticToHeat = transpose(HeatToElastic)
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)

        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)
        
        call elementHeatStiff3DT41(xyz, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass3DT41(xyz, coefM, tpHeatMass, HeatMass)
        call assembleMatrix(LHS, HeatMass + 0.5 * dt * HeatStiff, idxHeat)
        call assembleMatrix(RHS, HeatMass - 0.5 * dt * HeatStiff, idxHeat)
        
    end do

end subroutine

subroutine assembleDC2_T4CTE(this, rho, c, k, TRef, E, Nu, beta, dt, alpha, &
                             tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES3DT412CTE), intent(in) :: this
    real*8, intent(in):: rho, c, k, TRef, E, Nu, beta, dt, alpha
    integer, intent(in) :: tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled
    type(SMatrix), intent(inout) :: LHS, RHS  

    integer :: i, j, idx(4), idxHeat(4), idxElastic(12)
    real*8 :: xyz(4, 3), coefM(4), coefK(4), coefE(4), coefNu(4), coefBeta(4), coefP(4)
    real*8 :: HeatStiff(4, 4), HeatMass(4, 4)
    real*8 :: ElasticStiff(12, 12), ElasticMass(12, 12)
    real*8 :: HeatToElastic(12, 4), ElasticToHeat(4, 12)

    type(Mesh3DT4), pointer :: mesh3D
    mesh3D => this%mesh3D
    coefP = rho
    coefM = rho * c / TRef
    coefK = k / TRef
    coefE = E
    coefNu = Nu
    coefBeta = beta

    do i = 1, mesh3D%nt
        idx = mesh3D%elements(i)%idx
        xyz(:, 1) = mesh3D%nodes(idx)%xyz(1)
        xyz(:, 2) = mesh3D%nodes(idx)%xyz(2)
        xyz(:, 3) = mesh3D%nodes(idx)%xyz(3)
        call getDofFromElement(idx, 3, idxElastic)
        idxHeat = idx + 3 * mesh3D%nv

        call elementElasticityStiff3DT42(xyz, coefE, coefNu, tpElasticStiff, ElasticStiff)
        call elementElasticityMass3DT42(xyz, coefP, tpElasticMass, ElasticMass)
        call elementCoupledMatrixT4CTE(xyz, coefBeta, tpCoupled, HeatToElastic)
        ElasticToHeat = transpose(HeatToElastic)
        call elementHeatStiff3DT41(xyz, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass3DT41(xyz, coefM, tpHeatMass, HeatMass)

        call assembleMatrix(LHS, ElasticStiff + 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)     
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(LHS, HeatMass + 0.5 * dt * HeatStiff, idxHeat)

        call assembleMatrix(RHS, 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)
        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(RHS, HeatMass - 0.5 * dt * HeatStiff, idxHeat)    
        
    end do
end subroutine

subroutine assembleHeatConvectSurfaceT4CTE_Const(this, g, TRef, dt, label, tp, A, M)
    implicit none
    class(FES3DT412CTE), intent(in) :: this
    real*8, intent(in) :: g, TRef, dt
    type(smatrix), intent(inout) :: A, M
    integer, intent(in) :: label, tp

    integer :: i, j,idx(3)
    real*8 :: xyz(3, 3), gg(3), k(3, 3)
    type(Mesh3DT4), pointer :: mesh3D
    mesh3D => this%mesh3D
	gg = g/TRef
    do i = 1, mesh3D%nb
        if(mesh3D%boundarys(i)%num == label) then
            idx = mesh3D%boundarys(i)%idx
            xyz(:, 1) = mesh3D%nodes(idx)%xyz(1)
            xyz(:, 2) = mesh3D%nodes(idx)%xyz(2)
            xyz(:, 3) = mesh3D%nodes(idx)%xyz(3)
            idx = idx + 3 * mesh3D%nv
            call elementConvectSurface3DT41(xyz, gg, tp, k)
            call assembleMatrix(A, 0.5 * dt * K, idx)
            call assembleMatrix(M, -0.5 * dt * K, idx)
        end if
    end do    
end subroutine

subroutine assembleHeatFluxT4CTE_function(this, g, time, label, tp, RHS)
    implicit none
    class(FES3DT412CTE), intent(in) :: this         
    real*8, external :: g
    real*8, intent(in) :: time
    real*8, intent(inout) :: RHS(this%nDof)        
    integer, intent(in) :: label, tp

    integer :: i, j, idx(3)
    real*8 :: xyz(3, 3), gg(3), k(3)
    type(Mesh3DT4), pointer :: mesh3D
    mesh3D => this%mesh3D
    do i = 1, mesh3D%nb
        if(mesh3D%boundarys(i)%num == label) then
            idx = mesh3D%boundarys(i)%idx
            xyz(:, 1) = mesh3D%nodes(idx)%xyz(1)
            xyz(:, 2) = mesh3D%nodes(idx)%xyz(2)
            xyz(:, 3) = mesh3D%nodes(idx)%xyz(3)
            gg(1) = g(xyz(1, :), label, time)
            gg(2) = g(xyz(2, :), label, time)
            gg(3) = g(xyz(3, :), label, time)
            idx = idx + 3 * mesh3D%nv
            call elementHeatFlux3DT41(xyz, gg, tp, k)
            call assembleVector(RHS, k, idx)
        end if
    end do
end subroutine

subroutine assembleHeatFluxT4CTE_Const(this, g, label, tp, RHS)
    implicit none
    class(FES3DT412CTE), intent(in) :: this
    real*8, intent(in) :: g
    real*8, intent(inout) :: RHS(this%nDof)
    integer, intent(in) :: label, tp

    integer :: i, j, idx(3)
    real*8 :: xyz(3, 3), gg(3), k(3)
    type(Mesh3DT4), pointer :: mesh3D
    mesh3D => this%mesh3D
	gg = g
    do i = 1, mesh3D%nb
        if(mesh3D%boundarys(i)%num == label) then
            idx = mesh3D%boundarys(i)%idx
            xyz(:, 1) = mesh3D%nodes(idx)%xyz(1)
            xyz(:, 2) = mesh3D%nodes(idx)%xyz(2)
            xyz(:, 3) = mesh3D%nodes(idx)%xyz(3)
            idx = idx + 3 * mesh3D%nv
            call elementHeatFlux3DT41(xyz, gg, tp, k)
            call assembleVector(RHS, k, idx)
        end if
    end do
end subroutine


subroutine elementCoupledMatrixT4CTE(xyz, beta, tp, K)
    implicit none
    real*8, intent(in) :: xyz(4, 3), beta(4)
    integer, intent(in) :: tp
    real*8, intent(out) :: K(12, 4)

    real*8 :: coef(4), x(4), y(4), z(4)
    real*8 :: dx(4), dy(4), dz(4), volume, gradient(12)
    integer :: i, j
    x = xyz(:, 1); y = xyz(:, 2); z = xyz(:, 3)
    call gradient3DP1(x, y, z, dx, dy, dz, volume)
    coef = (sum(beta) + beta) * volume / 20
    gradient(1:12:3) = dx
    gradient(2:12:3) = dy
    gradient(3:12:3) = dz
    do i = 1, 12
        K(i, :) = gradient(i) * coef
    end do
end subroutine

subroutine extractUVWHT4CTE(this, x, U, V, W, H)
    implicit none
    class(FES3DT412CTE), intent(in) :: this
    real*8, intent(in) :: x(this%nDof)
    real*8, intent(out) :: H(this%nDof/4)
    real*8, intent(out) :: U(this%nDof/4), V(this%nDof/4), W(this%nDof/4)

    integer :: nv
    nv = this%nDof/4
    U = x(1:3*nv:3)
    V = x(2:3*nv:3)
    W = x(3:3*nv:3)
    H = x(3*nv+1:this%nDof)
end subroutine



subroutine applyBCZeroDispT4CTE(this, A, RHS, label, direct)
    implicit none
    class(FES3DT412CTE), intent(in) :: this
    type(SMatrix), intent(inout) :: A
    real*8, intent(inout) :: RHS(A%M)
    integer, intent(in) :: label(:)
    integer, intent(in), optional :: direct
    
    type(LinkedSet) :: dirichletBC
    integer, allocatable :: DOF(:)
    type(Mesh3DT4), pointer :: mesh3D
    mesh3D => this%mesh3D
    call mesh3D.extractBoundaryNodes(dirichletBC, label)
    if (present(direct)) then
        allocate(DOF(dirichletBC%size))
        call getDofFromSet(dirichletBC, 3, DOF, direct)
    else
        allocate(DOF(dirichletBC%size * 3))
        call getDofFromSet(dirichletBC, 3, DOF)
    end if
    call this.applyBCZeroDOF_MR(A, RHS, DOF)
    deallocate(DOF)
    call dirichletBC.empty()
end subroutine

subroutine applyBCHeatT4CTE_function(this, A, RHS, label, u0, time, tp)
    use Set_m
	implicit none
	class(FES3DT412CTE), intent(in) :: this
	type(SMatrix), intent(inout) :: A
	real*8, intent(inout) :: RHS(this%nDof)
	integer, intent(in) :: label(:)
	real*8, external :: u0
    real*8, intent(in) :: time
	integer, intent(in) :: tp

	real*8, allocatable :: uu0(:), BCDOF(:)
	integer :: i
	real*8 :: xyz(3)
	type(Mesh3DT4), pointer :: mesh3D
    type(LinkedSet) :: dirichletBC
	type(ele), pointer:: current_p
    mesh3D => this%mesh3D

	call mesh3D.extractBoundaryNodes(dirichletBC, label)

	allocate(uu0(dirichletBC%size))

	i = 0
    current_p => dirichletBC%head
    do while(associated(current_p))
		i = i + 1
		xyz = mesh3D%nodes(current_p%n)%xyz
		uu0(i) = u0(xyz, time)
        current_p%n = current_p%n + 3 * mesh3D%nv
        current_p => current_p%next
    end do
    
    call this.applyBC_SetDOF(A, RHS, dirichletBC, uu0, tp)
	deallocate(uu0)
	call dirichletBC.empty()
end subroutine

subroutine applyBCHeatT4CTE_const(this, A, RHS, label, u0, tp)
	implicit none
	class(FES3DT412CTE), intent(in) :: this
	type(SMatrix), intent(inout) :: A
	real*8, intent(inout) :: RHS(this%nDof)
	integer, intent(in) :: label(:)
	real*8, intent(in) :: u0
	integer, intent(in) :: tp

	real*8, allocatable :: uu0(:)
    type(LinkedSet) :: dirichletBC

    call this%mesh3D%extractBoundaryNodes(dirichletBC, label)
    call dirichletBC.addNumber(3 * this%mesh3D%nv)

    allocate(uu0(dirichletBC%size))
    uu0 = u0
	call this.applyBC_SetDOF(A, RHS, dirichletBC, uu0, tp)
	deallocate(uu0)
	call dirichletBC.empty()
end subroutine


end module